//============================================================================
// Name        : cg.c
// Author      : Yugi
// Version     : 1.0
// Description : 
//============================================================================

#include "cg.h"
#include "udefine.h"

/* 
* Generate a new call gragh.
* 
* RETURN:	new cg
* NOTES:	Return value should be freed by user
*/
PCG cg_new(){
	PCG pcg = (PCG) malloc(sizeof(CG));
	pcg->nodes = (PCG_NODE) malloc(DEFAULT_CG_SIZE * sizeof(CG_NODE));
	pcg->head = NULL;
	pcg->tail = NULL;
	pcg->size = 0;
	pcg->multiple = 1;
	pcg->stack = NULL;
	return pcg;
}

/* 
* Insert a basic block to the call gragh.
* 
* REQUIRE:	pcg not NULL
*			pc in text seg
* RETURN:	Index of the target basic block, if the basic block
*			exist, return the existing basic block index
* MODIFY:	pcg
*/
unsigned int cg_add_node(void *drcontext, PCG pcg, app_pc pc){
	unsigned int index;

	/* Return the index of node if existed */
	for (index = 0; index<pcg->size; index++)
		if (pcg->nodes[index].pc == pc)
			return index;

	pcg->nodes[pcg->size].pc = pc;
	pcg->nodes[pcg->size].pcfg = cfg_new();
	cfg_add_bb(drcontext, pcg->nodes[pcg->size].pcfg, pcg->nodes[pcg->size].pc);
	cfg_static_complete(drcontext, pcg->nodes[pcg->size].pcfg);

	pcg->size++;

	/* If pcg->size meets the max size, twice it */
	if (pcg->size == DEFAULT_CFG_SIZE * pcg->multiple){
		PCG_NODE nodes = pcg->nodes;

		pcg->multiple++;
		pcg->nodes = (PCG_NODE) malloc(DEFAULT_CG_SIZE * sizeof(CG_NODE)*pcg->multiple);

		for (index = 0; index<pcg->size; index++)
			pcg->nodes[index] = nodes[index];

		free(nodes);
	}

	return pcg->size-1;
}

/* 
* Insert an edge to the call gragh.
* 
* REQUIRE:	pcg not NULL
*			src & dst distinguishes, less than pcg->size
* RETURN:	true if edge added
* MODIFY:	pcg
*/
bool cg_add_edge(PCG pcg, unsigned int src, unsigned int dst){
	/* 
	* Iterate through pcfg->edges, note that pcfg->edges is represented
	* by pcfg->head & pcfg->tail
	*/
	PEDGE edge = pcg->head;
	
	while (edge){
		if (edge->src==src && edge->dst==dst)
			return false;

		edge = edge->next;
	}
	
	/* Add to pcfg->edges */
	edge = (PEDGE) malloc(sizeof(EDGE));
	edge->src = src;
	edge->dst = dst;
	edge->next = NULL;

	if (pcg->head){
		pcg->tail->next = edge;
		pcg->tail = edge;
	}else
		pcg->head = pcg->tail = edge;

	return true;
}

/* 
* Construct a call graph. 
* 
* REQUIRE:	pcg not NULL
*			module not NULL
*			bb not NULL
* MODIFY:	pcg
*/
void cg_parse(void *drcontext, instrlist_t *bb, module_data_t *module, PCG pcg){
	app_pc pc;				// basic block address
	instr_t *pLast;			// Last instruction of basic block
	PINT_STACK pStack;		// Call stack pointer

	/* Initial */
	pc = instr_get_app_pc(instrlist_first(bb));
	pLast = instrlist_last(bb);

	/* Only cares about app code */
	if (pc<(module->start) || pc>(module->end))
		return;

	/* Entry point of the app */
	if (pc == module->entry_point){
		pStack = (PINT_STACK) malloc(sizeof(INT_STACK));
		pStack->index = cg_add_node(drcontext, pcg, pc);
		pStack->previous = NULL;
		pcg->stack = pStack;
	}

	/* Empty call stack */
	if (!pcg->stack)
		return;

	switch (instr_get_opcode(pLast)){
		case OP_call:
			pc = instr_get_branch_target_pc(pLast);

			if (pc){
				unsigned int index = cg_add_node(drcontext, pcg, pc);
				cg_add_edge(pcg, pcg->stack->index, index);

				/* If it is new basic block, add to call stack */
				if (index == pcg->size-1){
					pStack = (PINT_STACK) malloc(sizeof(INT_STACK));
					pStack->index = index;
					pStack->previous = pcg->stack;
					pcg->stack = pStack;
				}
			}

			break;
		case OP_ret:
			/* Remove the top basic block on call stack */
			pStack = pcg->stack;
			pcg->stack = pcg->stack->previous;
			free(pStack);
			break;
		default:
			;
	}
}

/* 
* Print call gragh
* 
* REQUIRE:	pcg not NULL
*			outfile not NULL
*			path_buf not NULL
*			dir_len less than MAXIMUM_PATH
*/
void cg_print(void *drcontext, file_t outfile, PCG pcg, char *path_buf, unsigned int dir_len){
	PEDGE edge = pcg->head;	// Edge iterator
	unsigned int index;		// pcg index 
	unsigned int ret;		// return value
	file_t fCFG;			// file for cfg

	dr_fprintf(outfile, "digraph CG {\n");

	/* Iterate through pcg->edges */	
	while (edge){
		dr_fprintf(outfile, "\"cfg%d 0x%.8x\" -> \"cfg%d 0x%.8x\"\n", 
			edge->src, pcg->nodes[edge->src].pc, edge->dst, pcg->nodes[edge->dst].pc);
		edge = edge->next;
	}

	dr_fprintf(outfile, "}\n");

	/* Create directory */
	ret = dr_snprintf(path_buf+dir_len, MAXIMUM_PATH-dir_len, DIRCFG);

	if (ret <= 0){
		dr_messagebox("Client path too long, move it to ROOT");
		exit(PATH_TOO_LONG);
	}

	if (!dr_directory_exists(path_buf))
		dr_create_dir(path_buf);

	/* Print control flow graph for every ndoe in call graph */
	for (index=0; index<pcg->size; index++){
		/* Open output file */
		ret = dr_snprintf(path_buf+dir_len, MAXIMUM_PATH-dir_len, "%s/%s%d%s", DIRCFG, PREFIX, index, SUFFIX);

		if (ret <= 0){
			dr_messagebox("Client path too long, move it to ROOT");
			exit(PATH_TOO_LONG);
		}

		fCFG = dr_open_file(path_buf, DR_FILE_WRITE_OVERWRITE);

		/* Print control flow graph */
		pcg->nodes[index].pcfg = cfg_new();
		cfg_add_bb(drcontext, pcg->nodes[index].pcfg, pcg->nodes[index].pc);
		cfg_static_complete(drcontext, pcg->nodes[index].pcfg);
		cfg_print(drcontext, fCFG, pcg->nodes[index].pcfg, false);

		/* Flush & Close */
		dr_flush_file(fCFG);
		dr_close_file(fCFG);
	}
}
